import os
import shutil
import hogfeature
import svmclassify
import joblib

# 就在这里加载我们自定义的常量
import variable as va  

def mkdir():
    # 检查用于存放特征的文件夹是否存在，如果没有就创建
    if not os.path.exists(va.train_feat_path):
        os.mkdir(va.train_feat_path)
    if not os.path.exists(va.test_feat_path):
        os.mkdir(va.test_feat_path)

if __name__ == '__main__':
    clf = None
    # 询问是否重新提取特征
    need_reget = input('是否需要重新获取特征？y/n\n')
    if need_reget == 'y':
        # shutil.rmtree()函数能递归删除文件夹下的所有子文件夹和子文件。
        shutil.rmtree(va.train_feat_path)
        shutil.rmtree(va.test_feat_path)
        # 再重新创建用于存放特征的文件夹
        mkdir()
        # 获取特征
        hogfeature.get_feats()
        # 重新提取特征了就必须重新训练模型
        clf = svmclassify.train()
    if(clf == None):
        # 询问是否重新训练模型
        need_train = input('是否重新训练模型？y/n\n')
        if(need_train == 'y'):
            clf = svmclassify.train()
        else:
            clf = joblib.load(va.model_path+'model')
    # 进行测试
    svmclassify.test(clf)
